#include <mem.h>
#include <assert.h>
#include "rbtree.h"

rbt_node* _rbt_new_node( rbt_node *_parent );
void _rbt_destroy( rbt_node *_rbtn );
void _rbt_remove( rbtree *_rbt, rbt_node *_rbtn );
rbt_node* _rbt_right_most( rbt_node *_rbtn );
rbt_node* _rbt_left_most( rbt_node *_rbtn );
void _rbt_adjust_insert( rbtree *_rbt, rbt_node *_rbtn );
void _rbt_adjust_remove( rbtree *_rbt, rbt_node *_rbtn, int flag );
void _rbt_roll_ll( rbtree *_rbt, rbt_node *_rbtn );
void _rbt_roll_lr( rbtree *_rbt, rbt_node *_rbtn );
void _rbt_roll_rl( rbtree *_rbt, rbt_node *_rbtn );
void _rbt_roll_rr( rbtree *_rbt, rbt_node *_rbtn );
rbt_iterator _rbt_find( rbtree *_rbt, void *_key );

rbtree* rbt_create( compare _cmp )
{
	rbtree *_rbt;
	_rbt = (rbtree*) palloc( sizeof( rbtree ) );
	assert( _rbt != NULL );

	_rbt->cmp = _cmp;
	_rbt->size = 0;
	_rbt->root = _rbt_new_node( NULL );

	return _rbt;
}

void rbt_destroy( rbtree *_rbt )
{
	_rbt_destroy( _rbt->root );
	pfree( _rbt );
}

rbt_iterator rbt_insert( rbtree *_rbt, void *_key, void *_val )
{
	rbt_iterator it;
	it = _rbt_find( _rbt, _key );
	if( it->pl != NULL )
	{
		it->val = _val;
	}
	else
	{
		it->key = _key;
		it->val = _val;
		it->flag = 1;
		it->pl = _rbt_new_node( it );
		it->pr = _rbt_new_node( it );
		
		_rbt->root->flag = 0;
		_rbt_adjust_insert( _rbt, it );
	}
	return it;
}

void* rbt_remove( rbtree *_rbt, void *_key )
{
	void *data = NULL;
	rbt_node *_rbn = _rbt_find( _rbt, _key );
	if( _rbn->pl != NULL )
	{
		data = _rbn->val;
		_rbt_remove( _rbt, _rbn );
	}
	return data;
}

rbt_iterator rbt_find( rbtree *_rbt, void *_key )
{
	rbt_node *_rbtn;
	_rbtn = _rbt_find( _rbt, _key );
	if( _rbtn->pl == NULL )
		return rbt_end( _rbt );
	else
		return _rbtn;
}

rbt_iterator rbt_begin( rbtree *_rbt )
{
	return _rbt->root;
}

rbt_iterator rbt_end( rbtree *_rbt )
{
	rbt_node *_rbtn;
	_rbtn = _rbt->root;
	while( _rbtn->pl != NULL )
		_rbtn = _rbtn->pr;
	return _rbtn;
}

rbt_iterator rbt_next( rbt_iterator _it )
{
	rbt_node *_rbtn;
	_rbtn = _it;
	if( _rbtn->pr != NULL )
	{
		return _rbt_left_most( _rbtn->pr );
	}
	else
	{
		while( _rbtn->pp != NULL )
		{
			if( _rbtn->pp->pl == _rbtn )
				return _rbtn->pp;
			else
				_rbtn = _rbtn->pp;
		}
	}
	return _it->pr;
}

rbt_iterator rbt_pre( rbt_iterator _it )
{
	rbt_node *_rbtn;
	_rbtn = _it;
	if( _rbtn->pl != NULL )
	{
		return _rbt_right_most( _rbtn->pl );
	}
	else
	{
		while( _rbtn->pp != NULL )
		{
			if( _rbtn->pp->pr == _rbtn )
				return _rbtn->pp;
			else
				_rbtn = _rbtn->pp;
		}
	}
	return _it->pl;
}


rbt_node* _rbt_new_node( rbt_node *_parent )
{
	rbt_node *_rbtn;
	_rbtn = (rbt_node *) palloc( sizeof( rbt_node ) );
	assert( _rbtn != NULL );

	_rbtn->flag = 0;
	_rbtn->key = NULL;
	_rbtn->pl = NULL;
	_rbtn->pr = NULL;
	_rbtn->pp = _parent;
	_rbtn->val = NULL;

	return _rbtn;
}

void _rbt_destroy( rbt_node *_rbtn )
{
	if( _rbtn->pl != NULL )
		_rbt_destroy( _rbtn->pl );
	if( _rbtn->pr != NULL )
		_rbt_destroy( _rbtn->pr );
	pfree( _rbtn );
}

void _rbt_remove( rbtree *_rbt, rbt_node *_rbtn )
{
	if( _rbtn->pl->pl == NULL )
	{
		if( _rbtn->pp != NULL )
		{
			_rbtn->pr->pp = _rbtn->pp;
			if( _rbtn->pp->pl == _rbtn )
				_rbtn->pp->pl = _rbtn->pr;
			else
				_rbtn->pp->pr = _rbtn->pr;
		}
		else
		{
			_rbt->root = _rbtn->pr;
			_rbtn->pr->pp = NULL;
		}
	}
	else
	{
		rbt_node *_rm = _rbt_right_most( _rbtn->pl );
		_rbtn->key = _rm->key;
		_rbtn->val = _rm->val;
		if( _rm->pp->pl == _rm )
		{
			_rm->pl->pp = _rm->pp;
			_rm->pp->pl = _rm->pl;
		}
		else
		{
			_rm->pl->pp = _rm->pp;
			_rm->pp->pr = _rm->pl;
		}
		_rbt_adjust_remove( _rbt, _rm->pl, _rm->flag );
	}
}

rbt_node* _rbt_right_most( rbt_node *_rbtn )
{
	while( _rbtn->pr != NULL && _rbtn->pr->pl != NULL )
		_rbtn = _rbtn->pr;
	return _rbtn;
}

rbt_node* _rbt_left_most( rbt_node *_rbtn )
{
	while( _rbtn->pl != NULL && _rbtn->pl->pl != NULL )
		_rbtn = _rbtn->pl;
	return _rbtn;
}

void _rbt_adjust_remove( rbtree *_rbt, rbt_node *_rbtn, int flag )
{
	rbt_node *parent = _rbtn->pp;
	rbt_node *pl = parent->pl;
	rbt_node *pr = parent->pr;

	if( _rbtn->flag == 1 || flag == 1 )
	{
		_rbtn->flag = 0;
	}
	else if( _rbtn == pl && pr->flag == 1 )
	{
		_rbt_roll_rr( _rbt, pr );
	}
	else if( _rbtn == pr && pl->flag == 1 )
	{
		_rbt_roll_ll( _rbt, pl );
	}
	else if( _rbtn == pl && parent->flag == 1 && pr->flag == 0 )
	{
		parent->flag = 0;
		pr->flag = 1;
	}
	else if( _rbtn == pr && parent->flag == 1 && pl->flag == 0 )
	{
		parent->flag = 0;
		pl->flag = 1;
	}
	
}

void _rbt_adjust_insert( rbtree *_rbt, rbt_node *_rbtn )
{
	if( _rbtn->pp != NULL && _rbtn->pp->flag == 1)
	{
		rbt_node *_grand, *_parent;
		_parent = _rbtn->pp;
		_grand = _parent->pp;
		if( _grand->pl->flag == 1 && _grand->pr->flag == 1 )
		{
			_grand->pl->flag = 0;
			_grand->pr->flag = 0;
			_grand->flag = 1;
			_rbt_adjust_insert( _rbt, _grand );
		}
		else if( _parent->flag == 1 )
		{
			if( _grand->pl == _parent && _parent->pl == _rbtn )
				_rbt_roll_ll( _rbt, _rbtn );
			else if( _grand->pl == _parent && _parent->pr == _rbtn )
				_rbt_roll_lr( _rbt, _rbtn );
			else if( _grand->pr == _parent && _parent->pl == _rbtn )
				_rbt_roll_rl( _rbt, _rbtn );
			else if( _grand->pr == _parent && _parent->pr == _rbtn )
				_rbt_roll_rr( _rbt, _rbtn );
		}
	}
	_rbt->root->flag = 0;
}

void _rbt_roll_ll( rbtree *_rbt, rbt_node *_rbtn )
{
	rbt_node *_parent, *_grand, *_ggrand;
	_parent = _rbtn->pp;
	_grand = _parent->pp;
	_ggrand = _grand->pp;

	_grand->pl = _parent->pr;
	_grand->pp = _parent;
	_grand->flag = 1;

	_parent->pr->pp = _grand;

	_parent->pr = _grand;
	_parent->pp = _ggrand;
	_parent->flag = 0;
	if( _grand == _rbt->root )
		_rbt->root = _parent;
	else if( _ggrand->pl == _grand )
		_ggrand->pl = _parent;
	else
		_ggrand->pr = _parent;

}

void _rbt_roll_lr( rbtree *_rbt, rbt_node *_rbtn )
{
	rbt_node *_new, *_parent, *_grand, *_ggrand;
	_new = _rbtn;
	_parent = _new->pp;
	_grand = _parent->pp;
	_ggrand = _grand->pp;

	_parent->pr = _new->pl;
	_parent->pp = _new;

	_grand->pp = _new;
	_grand->pl = _new->pr;
	_grand->flag = 1;

	_new->pl->pp = _parent;
	_new->pr->pp = _grand;

	_new->pl = _parent;
	_new->pr = _grand;
	_new->pp = _ggrand;
	_new->flag = 0;
	if( _grand == _rbt->root )
		_rbt->root = _new;
	else if( _ggrand->pl == _grand )
		_ggrand->pl = _new;
	else
		_ggrand->pr = _new;
}

void _rbt_roll_rl( rbtree *_rbt, rbt_node *_rbtn )
{
	rbt_node *_new, *_parent, *_grand, *_ggrand;
	_new = _rbtn;
	_parent = _new->pp;
	_grand = _parent->pp;
	_ggrand = _grand->pp;

	_parent->pl = _new->pr;
	_parent->pp = _new;

	_grand->pp = _new;
	_grand->pr = _new->pl;
	_grand->flag = 1;

	_new->pl->pp = _grand;
	_new->pr->pp = _parent;

	_new->pl = _grand;
	_new->pr = _parent;
	_new->pp = _ggrand;
	_new->flag = 0;
	if( _grand == _rbt->root )
		_rbt->root = _new;
	else if( _ggrand->pl == _grand )
		_ggrand->pl = _new;
	else
		_ggrand->pr = _new;
}

void _rbt_roll_rr( rbtree *_rbt, rbt_node *_rbtn )
{
	rbt_node *_parent, *_grand, *_ggrand;
	_parent = _rbtn->pp;
	_grand = _parent->pp;
	_ggrand = _grand->pp;

	_grand->pr = _parent->pl;
	_grand->pp = _parent;
	_grand->flag = 1;

	_parent->pl->pp = _grand;

	_parent->pl = _grand;
	_parent->pp = _ggrand;
	_parent->flag = 0;
	if( _grand == _rbt->root )
		_rbt->root = _parent;
	else if( _ggrand->pl == _grand )
		_ggrand->pl = _parent;
	else
		_ggrand->pr = _parent;
}

rbt_iterator _rbt_find( rbtree *_rbt, void *_key )
{
	rbt_node *_rbtn;
	int c;
	_rbtn = _rbt->root;
	while( _rbtn->pl != NULL )
	{
		c = _rbt->cmp( _rbtn->key, _key );
		if( c == 0 )
			return _rbtn;
		else if( c < 0 )
			_rbtn = _rbtn->pr;
		else
			_rbtn = _rbtn->pl;
	}
	return _rbtn;
}
